#include <assert.h>
#include <stdlib.h>

#include <string>
#include <iostream>
#include <openssl/aes.h>

#include "crypto.hpp"
#include "misc.hpp"
#include "mallocSharedPtr.hpp"

namespace wheels
{

#ifndef OPENSSL_3X 
std::string encrypto( const uint8_t * data , size_t len , const char * key )
{
	if( !data ) std::runtime_error( "" );
	assert( key );

	std::string ret;
	AES_KEY aes;
	char tmpIV[] = "2624750004598718";
	uint8_t iv[ AES_BLOCK_SIZE ];

	for( int i = 0; i < 16; i ++ ){
		iv[ i ] = tmpIV[ i ];
	}

	int rst = AES_set_encrypt_key( (uint8_t*)key , 128 , &aes );
	if( rst < 0 ){
		throw std::runtime_error( " set aes key fail." );
	}

	uint8_t * __dat = ( uint8_t* )malloc( len * 3 + 1);
	if( __dat == nullptr ){
		throw std::runtime_error( " Allocate memory fail."  );
	}

	// 执行加密
	AES_cbc_encrypt( data  , __dat , len , &aes , iv , AES_ENCRYPT );
	char * __rst_data = ( char *)malloc( len * 3 + 1);
	if( __rst_data == nullptr ){
		free( __dat );
		__dat = nullptr;
		throw std::runtime_error( " Allocate memory fail."  );
	}
	for( size_t i = 0; i < len; i ++ ){
		sprintf( __rst_data + i * 2, "%02X", __dat[ i ] );
	}
	ret.assign( __rst_data );
	if( __rst_data ) free( __rst_data );
	if( __dat ) free( __dat );
	return ret;
}

std::string decrypto( const uint8_t * data , size_t len , const char * key )
{
	assert( data );
	assert( key );

	std::string ret;
	AES_KEY aes;
	char tmpIV[] = "2624750004598718";
	uint8_t iv[ AES_BLOCK_SIZE ];
	uint8_t* real_data;
	wheels::mallocSharedPtr<uint8_t> tmp_ptr( len );
	real_data = tmp_ptr.get();

	for( size_t i = 0; i < len; i += 2 ){
		unsigned char __tmp_data[ 3 ];
		__tmp_data[ 0 ] = data[ i ];
		__tmp_data[ 1 ] = data[ i + 1 ];
		__tmp_data[ 2 ] = 0;
		real_data[ i / 2 ] = ( uint8_t )strtoul( (char*)__tmp_data , nullptr , 16 );
	}
	for( int i = 0; i < 16; i ++ ){
		iv[ i ] = tmpIV[ i ];
	}

	int rst = AES_set_decrypt_key( (uint8_t*)key , 128 , &aes );
	if( rst < 0 ){
		throw std::runtime_error( "set aes key fail."  );
	}

	uint8_t * __dat = ( uint8_t* )malloc( len / 2 + 1 );

	if( __dat == nullptr ){
		throw std::runtime_error( "Allocate memory fail."  );
	}
	memset( __dat , 0 , len / 2 + 1 );
	AES_cbc_encrypt( real_data , __dat , len / 2 , &aes , iv , AES_DECRYPT );

	ret.assign( ( char *)__dat );
	free( __dat );
	return ret;
}

#else
#include <openssl/evp.h>
#include <openssl/rand.h>

std::string encrypto(const uint8_t* data, size_t len, const char* key) {
	if (!data)
		throw std::runtime_error("");

	if (!key)
		throw std::runtime_error("Key is null");

	std::string ret;
	std::array<uint8_t, AES_BLOCK_SIZE> iv;
	char tmpIV[] = "2624750004598718";

	for (int i = 0; i < AES_BLOCK_SIZE; i++) {
		iv[i] = tmpIV[i];
	}

	EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
	if (!ctx) {
		throw std::runtime_error("Failed to create cipher context");
	}

	if (EVP_EncryptInit_ex(ctx, EVP_aes_128_cbc(), nullptr, (const unsigned char*)key, iv.data()) != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Failed to initialize encryption");
	}

	size_t outlen;
	size_t ciphertext_len = len + AES_BLOCK_SIZE;
	std::vector<uint8_t> ciphertext(ciphertext_len);

	if (EVP_EncryptUpdate(ctx, ciphertext.data(), (int*)&outlen, data, len) != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Encryption failed");
	}

	ciphertext_len = outlen;

	if (EVP_EncryptFinal_ex(ctx, ciphertext.data() + outlen, (int*)&outlen) != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Encryption failed");
	}

	ciphertext_len += outlen;
	EVP_CIPHER_CTX_free(ctx);

	for (size_t i = 0; i < ciphertext_len; i++) {
		ret += std::to_string(ciphertext[i]);
	}

	return ret;
}



std::string decrypto(const uint8_t* data, size_t len, const char* key)
{
	assert(data);
	assert(key);

	std::string ret;
	AES_KEY aes;
	char tmpIV[] = "2624750004598718";
	uint8_t iv[AES_BLOCK_SIZE];
	uint8_t* real_data;
	std::unique_ptr<uint8_t[]> tmp_ptr(new uint8_t[len]);
	real_data = tmp_ptr.get();

	for (size_t i = 0; i < len; i += 2) {
		unsigned char __tmp_data[3];
		__tmp_data[0] = data[i];
		__tmp_data[1] = data[i + 1];
		__tmp_data[2] = 0;
		real_data[i / 2] = (uint8_t)strtoul((char*)__tmp_data, nullptr, 16);
	}
	for (int i = 0; i < 16; i++) {
		iv[i] = tmpIV[i];
	}

	EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
	if (ctx == nullptr) {
		throw std::runtime_error("Create cipher context failed.");
	}

	int rst = EVP_DecryptInit_ex(ctx, EVP_aes_128_cbc(), nullptr, (uint8_t*)key, iv);
	if (rst != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Set cipher key failed.");
	}

	int ciphertext_len = len / 2;
	// Calculate required buffer length for decrypted data
	int plaintext_len = ciphertext_len + AES_BLOCK_SIZE;
	std::unique_ptr<uint8_t[]> plaintext(new uint8_t[plaintext_len]);

	rst = EVP_DecryptUpdate(ctx, plaintext.get(), &plaintext_len, real_data, ciphertext_len);
	if (rst != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Decryption failed.");
	}

	int final_len = 0;
	rst = EVP_DecryptFinal_ex(ctx, plaintext.get() + plaintext_len, &final_len);
	if (rst != 1) {
		EVP_CIPHER_CTX_free(ctx);
		throw std::runtime_error("Decryption finalization failed.");
	}
	plaintext_len += final_len;

	ret.assign((char*)plaintext.get(), plaintext_len);
	EVP_CIPHER_CTX_free(ctx);
	return ret;
}
#endif
}

